# coding=utf8

import re
import os

# todo 将txt文件做成一列  一行一个字  回车保留 每一个非回车字符后加一个回车即可
# todo 将对应的ann 提取出来对应的竖排开始结尾 用IOBE标记
origin_data_dir = "ruijin_round1_train2_20181022/"
proc_data_dir = "train_round1_data_processed/"

file_list = os.listdir(origin_data_dir)

text_file_list = [x for x in file_list if x.endswith("txt")]

for textfile in text_file_list:
    with open(origin_data_dir + textfile, 'r') as rftxt:
        txt_str = rftxt.read()

    with open(origin_data_dir + textfile.split("txt")[0] + "ann", "r") as rfann:
        ann_str = rfann.read()
    # todo 由源文件字符串 统一做标记，计入词典 对某个位置应该标记的
    dict_ann = {}
    for line in ann_str.split("\n"):
        if len(line) == 0:
            continue
        try:
            entity_str = re.split(r"\t", line)[-1]
            attrib = re.split(r"\s", re.split(r"\t", line)[1])[0]
            indx_start = re.split(r"\s", re.split(r"\t", line)[1])[1]
            indx_end = re.split(r"\s", re.split(r"\t", line)[1])[-1]

            assert re.sub(r"\s|\n", "", txt_str[int(indx_start):int(indx_end)]) == re.sub(r"\n|\s", "", entity_str)
        except Exception as e:
            print("line goes hoo {}".format(line))
            continue
        for iinnd in range(int(indx_start), int(indx_end)):
            if iinnd == int(indx_start):
                dict_ann[iinnd] = "B-" + attrib
            elif iinnd == int(indx_end) - 1:
                dict_ann[iinnd] = "E-" + attrib
            else:
                dict_ann[iinnd] = "I-" + attrib

    proc_str = ""

    list_ind_dict_keys = dict_ann.keys()
    for ind in range(len(txt_str)):

        if ind not in list_ind_dict_keys:
            line2add = txt_str[ind] + " " + "o" + "\n"
        else:
            line2add = txt_str[ind] + " " + dict_ann[ind] + "\n"
        proc_str += line2add

    proc_str_single = re.sub(r"\n\n", "\n", proc_str)

    with open("/home/mm/PycharmProjects/ruijinyiyuan/train_round1_data_processed/"+textfile, "w") as wf:
        wf.write(proc_str_single)
    print("ko")
